from typing import Any

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
from pyspark.sql import DataFrame
from sparksampling.core.job.base_job import BaseJob
from sklearn.metrics import balanced_accuracy_score,accuracy_score,f1_score,precision_score


class RandomForestEvaluationJob(BaseJob):
    type_map = {
        'source_path': str,
        'selected_features_list': list,
        'key': str,
    }

    def __init__(self, source_path=None, selected_features_list=None, key=None, *args, **kwargs):
        super(RandomForestEvaluationJob, self).__init__(*args, **kwargs)
        self.source_path = source_path
        self.selected_features_list = selected_features_list
        self.key = key
        self.check_type()


    def _statistics(self, df: DataFrame, *args, **kwargs) -> dict:

        source_df = self._get_df_from_source(self.source_path, dataio=kwargs.get('data_io'))
        source_df_x = source_df.select(*self.selected_features_list).toPandas()
        source_df_y = source_df.select(*self.key).toPandas()
        df_x = df.select(*self.selected_features_list).toPandas()
        df_y=df.select(*self.key).toPandas()
        model = RandomForestClassifier()
        model.fit(df_x,df_y)
        prey_y = model.predict(source_df_x)
        report=classification_report(y_pred=prey_y,y_true=source_df_y)
        # score = score = balanced_accuracy_score(y_true=source_df_y,y_pred=prey_y)
        balanced_score = balanced_accuracy_score(y_true=source_df_y, y_pred=prey_y)
        ac_score = accuracy_score(y_true=source_df_y, y_pred=prey_y)
        f_score = (f1_score(y_true=source_df_y, y_pred=prey_y, pos_label='1')+f1_score(y_true=source_df_y, y_pred=prey_y, pos_label='0'))/2
        p_score = (precision_score(y_true=source_df_y, y_pred=prey_y, pos_label='1')+precision_score(y_true=source_df_y, y_pred=prey_y, pos_label='1'))/2
        return {
            "report":report,
            "balanced_score":balanced_score,
            "accuracy_score":ac_score,
            "f1_score":f_score,
            "precision_score":p_score
        }

